'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 2_study_step_back.py
* @Time: 2025/9/10
* @All Rights Reserve By Brtc
'''
import os

import dotenv
import weaviate
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.language_models import BaseLanguageModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
from langchain_core.retrievers import BaseRetriever
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_weaviate import WeaviateVectorStore
from weaviate.auth import AuthApiKey
dotenv.load_dotenv()
class StepBackRetriever(BaseRetriever):
    """回答回退检索器"""
    retriever:BaseRetriever
    llm:BaseLanguageModel
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> list[Document]:
        """根据传递的query执行问题回退并检索"""
        #1、构建少量提示词模板
        examples = [
            {"input":"博睿AI上有关于AI应用开发的课程吗？", "output":"博睿有哪些AI的课程"},
            {"input": "博小睿出生再那个国家？", "output": "博小睿的个人经历是怎样的？"},
            {"input": "司机可以开快车吗？", "output": "司机可以做什么？"}
        ]
        example_prompt= ChatPromptTemplate.from_messages([
            ("human","{input}"),
            ("ai", "{output}")
        ])

        few_shot_prompt = FewShotChatMessagePromptTemplate(
            examples = examples,
            example_prompt=example_prompt
        )
        # 2、构建生成会问题提示
        system_prompt = "你是一个世界知识的专家， 你的任务是回退问题，将问题改述更一般的前置问题，这样更容易回答，请严格参考实例来实现"
        prompt = ChatPromptTemplate.from_messages([
            ("system", system_prompt),
            few_shot_prompt,
            ("human","{question}")
        ])
        #3、构建生成回退问题链
        chain = (
            {"question":RunnablePassthrough()}
            |prompt
            |self.llm
            |StrOutputParser()
            |self.retriever
        )
        return chain.invoke(query)
# 导入数据
client = weaviate.connect_to_weaviate_cloud(
    skip_init_checks=True,
    cluster_url=os.getenv("WAEVIATE_URL"),
    auth_credentials=AuthApiKey(os.getenv("WEAVIATE_KEY"))
)
embedding = OpenAIEmbeddings(model="text-embedding-3-small")
db = WeaviateVectorStore(client=client,
                         index_name="DataSetTest",
                         text_key="text",
                         embedding=embedding)
#2、创建回答回退检索器
step_back_retriever = StepBackRetriever(
    retriever=db.as_retriever(),
    llm=ChatOpenAI(model="gpt-4o-mini", temperature=0)
)
#3、检索文档
docs = step_back_retriever.invoke("人工智能会让世界发生翻天地覆的变化吗？")
for doc in docs:
    print(doc)